import glob

from scipy import signal
import scipy.io.wavfile as wav
from fun_wav import *



def FiltHighWav(signalWav, N, highPass):
    b, a = signal.butter(N=N, Wn=highPass, btype='highpass')
    filtSignal = signal.filtfilt(b, a, signalWav)  # signalWav为要过滤的信号
    return filtSignal

def showWav(filePath):
    fs, signalWav = wav.read(filePath)

    curbefore = fs // 4
    curAfter = fs // 4
    signalWav = signalWav[curAfter:len(signalWav) - curAfter]

    signalWav = signalWav / max(abs(signalWav))

    t = np.linspace(0, len(signalWav) / fs, len(signalWav), endpoint=False)

    # filtSignal = FiltWav(signalWav=signalWav, N=2, lowPass=lowPass, highPass=highPass)
    filtSignal = signalWav
    # envelope = GetEnvelope(filtSignal)

    plt.xlabel('time(s)')
    plt.ylabel('amplitute')

    plt.title(filePath)
    plt.plot(t, filtSignal)
    plt.show()

def showFreq(filePath):
    data = SingleSTFT(filePath)
    t = data[0]
    f = data[1]
    Zxx = data[2]

    print(f'freq max: {np.max(np.abs(Zxx))}')

    # Zxx = SubAdjacent(Zxx)
    # Zxx = np.abs(Zxx)
    t = t[0:Zxx.shape[1]]

    Zxx = np.log10(Zxx)

    print(Zxx.shape)
    # plt.title(filePath)
    plt.xlabel('time(s)')
    plt.ylabel('freqency(hz)')

    plt.title(filePath)
    plt.pcolormesh(t, f, Zxx)
    plt.show()

def showAll(filePath):
    showWav(filePath)
    showFreq(filePath)

def loadFiltSaveWav(srcPath, dstPath):
    fs, curData = wav.read(srcPath)

    # lowpass = 2 * 19910 / fs
    # highpass = 2 * 20090 / fs
    # curData = FiltWav(curData, 2, lowpass, highpass)

    highPass = 2 * 16000 / fs
    curData = FiltHighWav(curData, 4, highPass)

    saveData = curData.astype('int16')

    wav.write(dstPath, fs, saveData)

def fromDirLoadFiltSaveWav(srcDir, dstDir, volunList):
    if not os.path.exists(dstDir):
        os.makedirs(dstDir)
    for curVolun in volunList:
        pattern = "*.wav"
        subSrcDir = os.path.join(srcDir, str(curVolun))
        searchPath = os.path.join(subSrcDir, pattern)
        allPath = []
        for path in glob.glob(searchPath):
            allPath.append(path)
        for srcPath in allPath:
            curName = srcPath.split(os.sep)[-1]
            print(srcPath, curName)
            subDstDir = os.path.join(dstDir, str(curVolun))
            if not os.path.exists(subDstDir):
                os.makedirs(subDstDir)
            dstPath = os.path.join(subDstDir, curName)
            loadFiltSaveWav(srcPath, dstPath)

if __name__ == '__main__':
    # fileName = 'convert_test/add_1.wav'
    # saveName = 'save_test.wav'
    # loadFiltSaveWav(fileName, saveName)
    # showAll(saveName)
    srcDir = 'F:\\dataset\\final_data\\ultra'
    dstDir = 'F:\\dataset\\final_data\\trans_ultra'
    volunList = [2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20]
    fromDirLoadFiltSaveWav(srcDir, dstDir, volunList)

